import random
import time

import numpy as np
import tensorflow as tf

from src.evaluator import has_precision_problem_occurred
from src.generator_utils import get_random_seed_tensor, exec_method_by_index

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"


def sample(p):
    p_mapped = [x*100 for x in p]
    total = sum(p_mapped)
    rand_num = random.uniform(0, total)
    while rand_num == 0:
        rand_num = random.uniform(0, total)
    accumulate = 0
    for i in range(len(p)):
        accumulate += p_mapped[i]
        if accumulate >= rand_num:
            return i

def sample_mutated(p, k=3):
    selected = []
    for t in range(k):
        maxnum = -1
        maxindex = -1
        for i in range(len(p)):
            if p[i] > maxnum and i not in selected:
                maxindex = i
                maxnum = p[i]
        selected.append(maxindex)
    p_selected = [p[i] for i in selected]
    p_mapped = [x*100 for x in p_selected]
    total = sum(p_mapped)
    rand_num = random.uniform(0, total)
    while rand_num == 0:
        rand_num = random.uniform(0, total)
    accumulate = 0
    for i in range(k):
        accumulate += p_mapped[i]
        if accumulate >= rand_num:
            return selected[i]


def init_p(iteration_times, b, r, operator, pattern, shape):
    p_tf = [0 for i in range(42)]
    p_torch = [0 for i in range(42)]
    p_mnn = [0 for i in range(42)]
    for t in range(iteration_times):
        print('iter {0}: sampling...'.format(t))
        # 初始化seed tensor
        tensor_seed = get_random_seed_tensor(shape)
        for index in range(42):
            # print('- executing mutation type {0}'.format(index + 1))
            new_tensor = exec_method_by_index(tensor_seed, index)
            # 如果该操作得到的内容是溢出的 说明效率不高 直接判定不成功
            if not isinstance(new_tensor, str):
                detect_res = has_precision_problem_occurred(new_tensor, operator, pattern, b, r)
                if detect_res[0]:
                    p_tf[index] += 1
                if detect_res[1]:
                    p_torch[index] += 1
                if detect_res[2]:
                    p_mnn[index] += 1
    res_tf = [x / iteration_times for x in p_tf]
    res_torch = [x / iteration_times for x in p_torch]
    res_mnn = [x / iteration_times for x in p_mnn]
    return res_tf, res_torch, res_mnn


def run_process(operator, pattern, b, r, shape):
    # p_matrix_tf, p_matrix_torch, p_matrix_mnn = init_p(500, b, r, operator, pattern, shape)
    # print(p_matrix_tf)
    # print(p_matrix_torch)
    # print(p_matrix_mnn)
    # assert 1==2

    # p_matrix_tf = \
    #     [0.046, 0.0, 0.0, 0.0, 0.064, 0.02, 0.016, 0.014, 0.03, 0.072, 0.108, 0.074, 0.078, 0.054, 0.05, 0.056, 0.08,
    #      0.008, 0.006, 0.074, 0.092, 0.09, 0.084, 0.082, 0.102, 0.076, 0.09, 0.06, 0.062, 0.09, 0.102, 0.066, 0.098,
    #      0.058, 0.08, 0.068, 0.068, 0.096, 0.082, 0.054, 0.062, 0.05]
    # p_matrix_torch = \
    #     [0.046, 0.0, 0.0, 0.0, 0.064, 0.02, 0.016, 0.014, 0.03, 0.072, 0.108, 0.074, 0.078, 0.054, 0.05, 0.056, 0.08,
    #      0.008, 0.006, 0.074, 0.092, 0.09, 0.084, 0.082, 0.102, 0.076, 0.092, 0.06, 0.062, 0.09, 0.102, 0.066, 0.098,
    #      0.058, 0.08, 0.068, 0.068, 0.096, 0.082, 0.054, 0.062, 0.05]
    # p_matrix_mnn = \
    #     [0.046, 0.0, 0.0, 0.0, 0.064, 0.02, 0.016, 0.014, 0.03, 0.072, 0.108, 0.074, 0.078, 0.054, 0.05, 0.056, 0.08,
    #      0.008, 0.006, 0.074, 0.092, 0.09, 0.084, 0.082, 0.102, 0.076, 0.092, 0.06, 0.062, 0.09, 0.102, 0.066, 0.098,
    #      0.058, 0.08, 0.068, 0.068, 0.096, 0.082, 0.054, 0.062, 0.05]
    # random
    # success_case_random_torch = 0
    # success_case_random_tf = 0
    # success_case_random_mnn = 0
    #
    # for t in range(15000):
    #     if t % 1000 == 0:
    #         print('{0} case.'.format(t + 1))
    #     seed_tensor = get_random_seed_tensor(shape)
    #     res = has_precision_problem_occurred(seed_tensor, operator, pattern, b, r)
    #     if res[0]:
    #         success_case_random_tf += 1
    #     if res[1]:
    #         success_case_random_torch += 1
    #     if res[2]:
    #         success_case_random_mnn += 1
    #
    # print('{0} data cases have successfully raised precision problems.--random-tf'.format(success_case_random_tf))
    # print('{0} data cases have successfully raised precision problems.--random-torch'.format(success_case_random_torch))
    # print('{0} data cases have successfully raised precision problems.--random-mnn'.format(success_case_random_mnn))

    # t = 0
    #
    # time_random = time.time()
    # while success_case_random_torch < 100 or success_case_random_tf < 100 or success_case_random_mnn < 100:
    #     if t % 100 == 0:
    #         print('{0} case.'.format(t + 1))
    #     seed_tensor = get_random_seed_tensor(shape)
    #     res = has_precision_problem_occurred(seed_tensor, operator, pattern, b, r)
    #     if res[0]:
    #         success_case_random_tf += 1
    #     seed_tensor = get_random_seed_tensor(shape)
    #     res = has_precision_problem_occurred(seed_tensor, operator, pattern, b, r)
    #     if res[1]:
    #         success_case_random_torch += 1
    #     seed_tensor = get_random_seed_tensor(shape)
    #     res = has_precision_problem_occurred(seed_tensor, operator, pattern, b, r)
    #     if res[2]:
    #         success_case_random_mnn += 1
    #     t += 1
    #     # print('---------')
    #     # print(success_case_random_tf)
    #     # print(success_case_random_torch)
    #     # print(success_case_random_mnn)
    #     # print('---------')
    # time_random = time.time() - time_random
    #
    # print('random: {0}'.format(time_random))
    # print(t)
    # print('{0} data cases have successfully raised precision problems.--random-tf'.format(success_case_random_tf))
    # print('{0} data cases have successfully raised precision problems.--random-torch'.format(success_case_random_torch))
    # print('{0} data cases have successfully raised precision problems.--random-mnn'.format(success_case_random_mnn))

    # tf torch mnn
    # success_case_tf = 0
    # success_case_torch = 0
    # success_case_mnn = 0

    success_case_tf_better = 0
    success_case_torch_better = 0
    success_case_mnn_better = 0
    t = 0
    time_weighted = time.time()
    while time.time() - time_weighted <= int(sys.argv[1]):
    # for t in range(15000):
        if t % 1000 == 0:
            print('{0} case.'.format(t + 1))
        # 分别采样mutation methods
        if sys.argv[2] == 'enforce':
            k = int(sys.argv[3])
            index_sampled_tf = sample_mutated(p_matrix_tf, k)
            index_sampled_torch = sample_mutated(p_matrix_torch, k)
            index_sampled_mnn = sample_mutated(p_matrix_mnn, k)       
        else:
            index_sampled_tf = sample(p_matrix_tf)
            index_sampled_torch = sample(p_matrix_torch)
            index_sampled_mnn = sample(p_matrix_mnn)
        seed_tensor = get_random_seed_tensor(shape)

        '''index_sampled_tf_b = sample_mutated(p_matrix_tf, 1)
        final_tensor_tf_b = exec_method_by_index(seed_tensor, index_sampled_tf_b)
        if not has_precision_problem_occurred(seed_tensor, operator, pattern, b, r)[0] && has_precision_problem_occurred(final_tensor_tf_b, operator, pattern, b, r)[0]:
            np.save('./data/tf/{}/{}'.format(operator, success_case_tf_better), seed_tensor)
            np.save('./data/tf/{}/mutated_{}'.format(operator, success_case_tf_better), final_tensor_tf_b)
            success_case_tf_better += 1

        index_sampled_torch_b = sample_mutated(p_matrix_torch, 1)
        final_tensor_torch_b = exec_method_by_index(seed_tensor, index_sampled_torch_b)
        if has_precision_problem_occurred(final_tensor_torch_b, operator, pattern, b, r)[1]:
        # np.save('./data/torch/{}_seed_{}'.format(operator, success_case_torch_better), seed_tensor)
            np.save('./data/torch/{}/mutated_{}'.format(operator, success_case_torch_better), final_tensor_torch_b)
            success_case_torch_better += 1
                
        index_sampled_mnn_b = sample_mutated(p_matrix_mnn, 1)
        final_tensor_mnn_b = exec_method_by_index(seed_tensor, index_sampled_mnn_b)
        if has_precision_problem_occurred(final_tensor_mnn_b, operator, pattern, b, r)[2]:
        # np.save('./data/mnn/{}_seed_{}'.format(operator, success_case_mnn_better), seed_tensor)
            np.save('./data/mnn/{}/mutated_{}'.format(operator, success_case_mnn_better), final_tensor_mnn_b)
            success_case_mnn_better += 1'''
            
        # print(success_case_mnn_better)
        # print(success_case_torch_better)
        # print(success_case_tf_better)

        # 生成新的tensor
        final_tensor_tf = exec_method_by_index(seed_tensor, index_sampled_tf)
        final_tensor_torch = exec_method_by_index(seed_tensor, index_sampled_torch)
        final_tensor_mnn = exec_method_by_index(seed_tensor, index_sampled_mnn)



        # 更新
        if has_precision_problem_occurred(final_tensor_tf, operator, pattern, b, r)[0]:
            success_case_tf += 1

        if has_precision_problem_occurred(final_tensor_torch, operator, pattern, b, r)[1]:
            success_case_torch += 1

        if has_precision_problem_occurred(final_tensor_mnn, operator, pattern, b, r)[2]:
            success_case_mnn += 1

        t += 1

    time_weighted = time.time() - time_weighted

    print('weighted: {0}'.format(time_weighted))
    print(t)

    # print('{0} data cases have successfully raised precision problems.--tf'.format(success_case_tf))
    # print('{0} data cases have successfully raised precision problems.--torch'.format(success_case_torch))
    # print('{0} data cases have successfully raised precision problems.--mnn'.format(success_case_mnn))
    print(success_case_tf_better)
    print(success_case_torch_better)
    print(success_case_mnn_better)


if __name__ == '__main__':
    # # batch_normalization层需要开启训练状态 使用样本均值和方差而不是默认均值或方差
    operator = sys.argv[4]
    if operator == 'batch_normalization':
        tf.keras.backend.set_learning_phase(1)
    pattern = sys.argv[5]
    b = float(sys.argv[6])
    r = float(sys.argv[6])
    shape = tuple(sys.argv[7])

    run_process(operator, pattern, b, r, shape)

